package com.qiangesoft.ratelimit.core.sentinel;

import com.alibaba.csp.sentinel.Entry;
import com.alibaba.csp.sentinel.SphU;
import com.alibaba.csp.sentinel.slots.block.BlockException;
import com.alibaba.csp.sentinel.slots.block.RuleConstant;
import com.alibaba.csp.sentinel.slots.block.flow.FlowRule;
import com.alibaba.csp.sentinel.slots.block.flow.FlowRuleManager;
import com.qiangesoft.ratelimit.core.LimitKeyBuilder;
import com.qiangesoft.ratelimit.core.LimitType;
import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.core.annotation.AnnotationUtils;
import org.springframework.stereotype.Component;

import java.util.ArrayList;
import java.util.List;
import java.util.Objects;

/**
 * sentinel限流切面
 *
 * @author qiangesoft
 * @date 2024-03-21
 */
@Aspect
@Component
public class SentinelLimitAspect {

    @Pointcut(value = "@annotation(com.qiangesoft.ratelimit.core.sentinel.SentinelRateLimit)")
    public void pointcut() {
    }

    @Around(value = "pointcut()")
    public Object around(ProceedingJoinPoint joinPoint) throws Throwable {
        // 1.从方法获取注解
        SentinelRateLimit sentinelRateLimit = this.getAnnotation(joinPoint);
        if (Objects.isNull(sentinelRateLimit)) {
            return joinPoint.proceed();
        }

        // 2.限流逻辑
        LimitType limitType = sentinelRateLimit.limitType();
        String key = sentinelRateLimit.key();
        int count = sentinelRateLimit.count();
        String limitKey = LimitKeyBuilder.build(limitType, key, joinPoint);
        initFlowRule(limitKey, count);
        Entry entry = null;
        try {
            entry = SphU.entry(limitKey);
            return joinPoint.proceed();
        } catch (BlockException ex) {
            throw new RuntimeException("访问太频繁，请稍后再试！");
        } finally {
            if (entry != null) {
                entry.exit();
            }
        }
    }

    /**
     * 获取注解
     *
     * @param joinPoint
     * @return
     */
    private SentinelRateLimit getAnnotation(JoinPoint joinPoint) {
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        return AnnotationUtils.findAnnotation(signature.getMethod(), SentinelRateLimit.class);
    }

    /**
     * 设置规则
     *
     * @param resourceName
     * @param limitCount
     */
    private static void initFlowRule(String resourceName, int limitCount) {
        List<FlowRule> rules = new ArrayList<>();
        FlowRule rule = new FlowRule();
        rule.setResource(resourceName);
        // 设置流控规则 QPS
        rule.setGrade(RuleConstant.FLOW_GRADE_QPS);
        // 设置受保护的资源阈值
        rule.setCount(limitCount);
        rules.add(rule);
        FlowRuleManager.loadRules(rules);
    }
}